from torch import nn, Tensor, optim
from torch.utils.data import Dataset, DataLoader
import torch

# 执行代码文件，会训练模型，并打印训练情况。

# 定义数据集

# 自定义数据集类
# 数据集中有以下样本：
# | 样本数据     | 样本标签 |
# | :----        | :----    |
# | [0.01, 0.02, 0.03] | [ 0.1 ]    |
# | [0.04, 0.05, 0.06] | [ 0.2 ]    |
# | [0.07, 0.08, 0.09] | [ 0.3 ]    |
# | [0.10, 0.11, 0.12] | [ 0.4 ]    |
# | [0.13, 0.14, 0.15] | [ 0.5 ]    |
class CustomDataset(Dataset):

    # 初始化自定义数据集
    def __init__(self, transform=None, target_transform=None):
        self.transform = transform
        self.target_transform = target_transform

    # __len__方法需要返回数据集中样本的数量
    def __len__(self):
        return 5

    # 数据集在被使用的时候，__getitem__方法会被调用。
    # 参数idx是索引，对应__getitem__需要返回的第idx个样本的数据和对应的标签
    def __getitem__(self, idx):
        dataset = [
            {'data': [0.01, 0.02, 0.03], 'label': [ 0.1 ]},
            {'data': [0.04, 0.05, 0.06], 'label': [ 0.2 ]},
            {'data': [0.07, 0.08, 0.09], 'label': [ 0.3 ]},
            {'data': [0.10, 0.11, 0.12], 'label': [ 0.4 ]},
            {'data': [0.13, 0.14, 0.15], 'label': [ 0.5 ]}
        ]
        data = dataset[idx]['data']
        label = dataset[idx]['label']
        if self.transform:
            data = self.transform(data)
        if self.target_transform:
            label = self.target_transform(label)
        return data, label

# 定义神经网络模型

class NeuralNetwork(nn.Module):

    def __init__(self):
        super().__init__()
        # 这里构造了一个模型
        self.sequential = nn.Sequential(
            nn.Linear(3, 5),
            nn.Tanh(),
            nn.Linear(5, 1),
            nn.Tanh()
        )

    # 执行模型的方法，不用手动调用，是框架内部会调用
    def forward(self, x):
        return self.sequential(x)

# 实例化处理数据集和加载数据集

batch_size = 16
test_data = CustomDataset(Tensor, Tensor)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 实例化模型

model = NeuralNetwork()

# 打印模型结构

print('model:')
print(model)

# 准备训练所需的内容

loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.3)

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        loss, current = loss.item(), (batch + 1) * len(X)
        print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

def test(dataloader, model, loss_fn):
    model.eval()
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            print(pred)
            print(y)

epochs = 10000
for t in range(epochs):
    print(f"Epoch {t+1}")
    train(test_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
